package com.woniuxy.vector;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.web.bind.annotation.*;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @Author: 马宇航
 * @Todo: 本地embedding向量存储，接口调试
 * @DateTime: 25/09/09/星期二 16:51
 * @Component: 成都蜗牛学苑
 **/
@RestController
@RequestMapping("/localVector")
public class LocalVectorController {
    private static final Logger logger = LoggerFactory.getLogger(LocalVectorController.class);
    private String LOCAL_PATH = "D:\\Desktop\\115期\\三阶段课堂项目\\new\\git-java-115\\myprojectAi\\vectorstore.json";
    private SimpleVectorStore vectorStore;
    @Autowired
    private ChatClient client;

    public LocalVectorController(EmbeddingModel embeddingModel) {
        logger.info("初始化本地向量存储");
        this.vectorStore = SimpleVectorStore.builder(embeddingModel).build();

    }
    /**
     * 本地向量存储add
     * @param doc
     * @return
     */
    @GetMapping("/add")
    public String add(@RequestParam String doc) throws IOException {
        //1.文档内容，通过代码处理
        String[] split = doc.split("。");
        List<Document> documents = new ArrayList<>();
        for (String s : split) {
            //后面的元数据，是用来快速区分文档的。
            documents.add(new Document(s, Map.of("name","蜗牛学苑","content","智能客服")));
        }
        vectorStore.add(documents);
//        try {
//            // 需要指定本地向量存储的路径 文件名
//            ClassPathResource classPathResource = new ClassPathResource("vectorstore.json");
//            //获取这个文件名所在的目录地址，赋值遍历，后续要用。
//            LOCAL_PATH = classPathResource.getFile().getAbsolutePath();
            vectorStore.save(new File(LOCAL_PATH));
//        } catch (IOException e) {
//            e.printStackTrace();
//            //创建文件
//            Files.createFile(Paths.get(LOCAL_PATH));
//            throw new RuntimeException(e);
//        }
        return "success";
    }
    @GetMapping("/search")
    public String search(@RequestParam String ask){
        //1. 从文档中search, 把文档的向量加载到内存中
        vectorStore.load(new File(LOCAL_PATH));
        //2. 从内存中搜索
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder()
                        .query(ask)
                        .topK(2) //返回2条
                        .similarityThreshold(0.5) //相似度阈值
                .build());
        return client.prompt(documents.toString()).user(ask).call().content();
    }
}
